package ai.koog.agents.example.banking.routing

import ai.koog.agents.core.agent.AIAgent
import ai.koog.agents.core.agent.AIAgentService
import ai.koog.agents.core.agent.createAgentTool
import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.core.tools.reflect.asTools
import ai.koog.agents.example.ApiKeyService
import ai.koog.agents.example.banking.tools.MoneyTransferTools
import ai.koog.agents.example.banking.tools.TransactionAnalysisTools
import ai.koog.agents.example.banking.tools.bankingAssistantSystemPrompt
import ai.koog.agents.example.banking.tools.transactionAnalysisPrompt
import ai.koog.agents.ext.tool.AskUser
import ai.koog.prompt.executor.clients.openai.OpenAIModels
import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
import kotlinx.coroutines.runBlocking

fun main() = runBlocking {
    simpleOpenAIExecutor(ApiKeyService.openAIApiKey).use { executor ->
        val transferAgentService = AIAgentService(
            promptExecutor = executor,
            llmModel = OpenAIModels.Chat.GPT4oMini,
            systemPrompt = bankingAssistantSystemPrompt,
            temperature = 0.0,
            toolRegistry = ToolRegistry { tools(MoneyTransferTools().asTools()) }
        )

        val analysisAgentService = AIAgentService(
            promptExecutor = executor,
            llmModel = OpenAIModels.Chat.GPT4oMini,
            systemPrompt = bankingAssistantSystemPrompt + transactionAnalysisPrompt,
            temperature = 0.0,
            toolRegistry = ToolRegistry { tools(TransactionAnalysisTools().asTools()) }
        )

        val classifierAgent = AIAgent(
            promptExecutor = executor,
            llmModel = OpenAIModels.Chat.GPT4oMini,
            toolRegistry = ToolRegistry {
                tool(AskUser)
                tool(
                    transferAgentService.createAgentTool(
                        agentName = "transferMoney",
                        agentDescription = "Transfers money and solves all arising problems",
                        inputDescription = "Transfer request"
                    )
                )
                tool(
                    analysisAgentService.createAgentTool(
                        agentName = "analyzeTransactions",
                        agentDescription = "Performs analytics for user transactions",
                        inputDescription = "Transaction analytics request"
                    )
                )
            },
            systemPrompt = bankingAssistantSystemPrompt + transactionAnalysisPrompt
        )

        println("Banking Assistant started")
        val message = "Send 25 euros to Daniel for dinner at the restaurant."
        // transfer messages
//        "Send 50 euros to Alice for the concert tickets"
//        "What's my current balance?"
        // analysis messages
//        "How much have I spent on restaurants this month?"
//         "What's my maximum check at a restaurant this month?"
//         "How much did I spend on groceries in the first week of May?"
//         "What's my total spending on entertainment in May?"
        val result = classifierAgent.run(message)
        println(result)
    }
}
